import { CCCAgentsError } from '../CCCAgentsError';
import { Utilities } from '../Utilities';
import { IActuator } from './IActuator';
import { IDiscreteActionMask } from './IDiscreteActionMask';

/**
 * 实现IDiscreteActionMask，允许从IActuator写入动作掩码
 */
export class ActuatorDiscreteActionMask implements IDiscreteActionMask {
	/**
	 * 使用离散控制（discrete control）时，为动作的起始指标。
	 */
	protected _startingActionIndices: Array<number> = null!;

	protected _branchSizes?: Array<number>;

	protected _currentMask: Array<boolean> = null!;

	protected _actuators: Array<IActuator> = null!;

	readonly _sumOfDiscreteBranchSizes: number;

	readonly _numBranches: number;

	currentBranchOffset: number = null!;

	setActionEnable(branch: number, actionIndex: number, isEnabled: boolean): void {
		this.lazeInitialize();

		const _actionIndex = this._branchSizes?.at(this.currentBranchOffset + branch) ?? 0;

		if (branch >= this._numBranches || actionIndex >= _actionIndex) {
			throw new RangeError('Invalid Action Masking: Action Mask is too large for specified branch.');
		}

		this._currentMask[actionIndex + this._startingActionIndices[this.currentBranchOffset + branch]] = !isEnabled;
	}

	lazeInitialize() {
		if (!this._branchSizes) {
			this._branchSizes = new Array(this._numBranches);

			this._actuators.forEach((e) => {
				this._branchSizes = this._branchSizes?.concat(e.actionSpec.branchSizes);
			});
		}

		if (!this._currentMask) {
			this._currentMask = new Array(this._sumOfDiscreteBranchSizes).fill(true);
		}

		if (!this._startingActionIndices) {
			this._startingActionIndices = Utilities.cumSum(this._branchSizes);
		}
	}

	constructor(
		actuator: Array<IActuator>,
		sumOfDiscreteBranchSizes: number,
		numBranches: number,
		branchSizes?: number[]
	) {
		this._actuators = actuator;
		this._sumOfDiscreteBranchSizes = sumOfDiscreteBranchSizes;
		this._numBranches = numBranches;
		this._branchSizes = branchSizes;
	}

	//todo
	getMask(): boolean[] {
		if (this._currentMask) {
			this.assertMask();
		}

		return this._currentMask;
	}

	assertMask() {
		for (let index = 0; index < this._numBranches; index++) {
			if (this.areAllActionsMasked(index)) {
				throw new CCCAgentsError(`Invalid Action Masking: All the actions of branch
                    ${index} are masked.
                `);
			}
		}
	}

	/**
	 * 重置当前agent掩码
	 */
	resetMask() {
		if (this._currentMask != null) {
			this._currentMask = new Array<boolean>(this._currentMask.length);
		}
	}

	/**
	 * 检查输入分支中的所有操作是否被屏蔽
	 * @param index 检测的分支
	 * @returns
	 */
	areAllActionsMasked(index: number): boolean {
		if (this._currentMask == null) {
			return false;
		}

		const start = this._startingActionIndices[index];
		const end = this._startingActionIndices[index + 1];
		for (let i = start; i < end; i++) {
			if (!this._currentMask[i]) {
				return false;
			}
		}

		return true;
	}
}
